import datetime
import numpy as np
import pandas as pd
import folium
import webbrowser
from folium.plugins import HeatMap
from matplotlib import pyplot as plt


def draw(data, path, name):
    num = len(data)
    lat = np.array(data["EPI_LAT"][0:num], dtype=float)  # 获取维度之维度值
    lon = np.array(data["EPI_LON"][0:num], dtype=float)  # 获取经度值
    m = np.array(data["M"][0:num], dtype=float)  # 获取人口数，转化为numpy浮点型

    data1 = [[lat[i], lon[i], m[i]] for i in range(num)]  # 将数据制作成[lats,lons,M]的形式

    map_osm = folium.Map(location=[31, 121], tiles='Stamen Terrain', zoom_start=2)  # 绘制Map，开始缩放程度是5倍
    HeatMap(data1).add_to(map_osm)  # 将热力图添加到前面建立的map里
    map_osm.save(path + name)
    webbrowser.open(path + name)


if __name__ == "__main__":
    df = pd.read_excel("全球地震数据.xlsx")
    begin_time = datetime.datetime(2019, 4, 17)
    end_time = datetime.datetime.now()
    df["is_in_range"] = df["O_TIME"].map(
        lambda x: (datetime.datetime.strptime(x, "%Y-%m-%d %H:%M:%S") >= begin_time) &
                  (datetime.datetime.strptime(x, "%Y-%m-%d %H:%M:%S") <= end_time)
    )
    df["datetime"] = df["O_TIME"].map(lambda x: datetime.datetime.strptime(x, "%Y-%m-%d %H:%M:%S"))
    df["year"] = df["datetime"].map(lambda x: x.year)
    df["month"] = df["datetime"].map(lambda x: x.month)
    df["year_month"] = df["year"] *12 + df["month"]
    draw(df.groupby("is_in_range").get_group(True), "./", "近一年全球地震分布.html")

    plt.rcParams['font.sans-serif'] = ['SimHei']

    fig1 = plt.figure()
    plt.plot(df["year"].value_counts().sort_index(), "r+-")
    plt.title("近年来全球地震次数分布")
    plt.xlabel("年份")
    plt.ylabel("次数")
    plt.savefig("近年来全球地震次数.png")
    plt.close()

    fig2 = plt.figure()
    plt.plot(df[df["M"] > 5]["year"].value_counts().sort_index(), "r+-")
    plt.title("五级以上近年来全球地震次数")
    plt.xlabel("年份")
    plt.ylabel("次数")
    plt.savefig("五级以上近年来全球地震次数.png")
    plt.close()

    fig3 = plt.figure()
    static = df[["year", "month",'M']].groupby(["year", "month"]).count().reset_index()
    static["year_month"] = static["year"].map(str) + "_" + static["month"].map(str)
    fig, ax = plt.subplots(figsize=(30, 15))
    plt.plot(static["year_month"], static["M"], "r+-")
    ax.set_xticklabels(list(static["year_month"]),
                       fontsize=16,
                       rotation=90)
    plt.tight_layout()
    plt.title("近年来全球地震次数（按月）")
    plt.xlabel("月份")
    plt.ylabel("次数")
    plt.savefig("近年来全球地震次数_按月.png")
    plt.close()

    fig3 = plt.figure()
    static = df[df["M"] > 5][["year", "month", 'M']].groupby(["year", "month"]).count().reset_index()
    static["year_month"] = static["year"].map(str) + "_" + static["month"].map(str)
    fig, ax = plt.subplots(figsize=(30, 15))
    plt.plot(static["year_month"], static["M"], "r+-")
    ax.set_xticklabels(list(static["year_month"]),
                       fontsize=16,
                       rotation=90)
    plt.tight_layout()
    plt.title("五级以上近年来全球地震次数（按月）")
    plt.xlabel("月份")
    plt.ylabel("次数")
    plt.savefig("五级以上近年来全球地震次数_按月.png")
    plt.close()